﻿using BuddHa.Common.SMUtil.Implementation.V1;
using Org.BouncyCastle.Crypto;
using Org.BouncyCastle.Crypto.Digests;
using Org.BouncyCastle.Crypto.Parameters;
using Org.BouncyCastle.Math.EC;
using Org.BouncyCastle.Security;
using Org.BouncyCastle.Utilities;
using System;

namespace BuddHa.Common.SMUtil.Implementation.V2
{
    internal class SM2Decrypter
    {
        private readonly IDigest _mDigest;
        private ECKeyParameters _mECKey;
        private ECDomainParameters _mECParams;
        private int _mCurveLength;

        public SM2Decrypter()
          : this(new SM3Digest())
        {
        }

        public SM2Decrypter(IDigest digest)
        {
            _mDigest = digest;
        }

        public virtual void Init(ICipherParameters param)
        {
            _mECKey = (ECKeyParameters)param;
            _mECParams = _mECKey.Parameters;
            _mCurveLength = (_mECParams.Curve.FieldSize + 7) / 8;
        }

        public byte[] Decrypt(byte[] data)
        {
            return Decrypt(data, 0, data.Length);
        }

        private byte[] Decrypt(byte[] input, int inOff, int inLen)
        {
            _mDigest.Reset();
            byte[] c1 = new byte[_mCurveLength * 2 + 1];

            Array.Copy(input, inOff, c1, 0, c1.Length);

            var c1P = _mECParams.Curve.DecodePoint(c1);
            var s = c1P.Multiply(_mECParams.H);
            if (s.IsInfinity)
                throw new InvalidCipherTextException("[h]C1 at infinity");

            c1P = c1P.Multiply(((ECPrivateKeyParameters)_mECKey).D).Normalize();

            byte[] c2 = new byte[inLen - c1.Length - _mDigest.GetDigestSize()];

            Array.Copy(input, inOff + c1.Length, c2, 0, c2.Length);

            Kdf(_mDigest, c1P, c2);

            AddFieldElement(_mDigest, c1P.AffineXCoord);
            _mDigest.BlockUpdate(c2, 0, c2.Length);
            AddFieldElement(_mDigest, c1P.AffineYCoord);

            byte[] c3 = DigestUtilities.DoFinal(_mDigest);

            int check = 0;
            for (int i = 0; i != c3.Length; i++)
            {
                check |= c3[i] ^ input[inOff + c1.Length + c2.Length + i];
            }

            Arrays.Fill(c1, 0);
            Arrays.Fill(c3, 0);

            if (check != 0)
            {
                Arrays.Fill(c2, 0);
                throw new InvalidCipherTextException("invalid cipher text");
            }

            return c2;
        }

        private void AddFieldElement(IDigest digest, ECFieldElement v)
        {
            byte[] p = v.GetEncoded();
            digest.BlockUpdate(p, 0, p.Length);
        }

        private void Kdf(IDigest digest, ECPoint c1, byte[] encData)
        {
            int digestSize = digest.GetDigestSize();
            byte[] buf = new byte[Math.Max(4, digestSize)];
            int off = 0;

            IMemoable memo = digest as IMemoable;
            IMemoable copy = null;

            if (memo != null)
            {
                AddFieldElement(digest, c1.AffineXCoord);
                AddFieldElement(digest, c1.AffineYCoord);
                copy = memo.Copy();
            }

            uint ct = 0;

            while (off < encData.Length)
            {
                if (memo != null)
                {
                    memo.Reset(copy);
                }
                else
                {
                    AddFieldElement(digest, c1.AffineXCoord);
                    AddFieldElement(digest, c1.AffineYCoord);
                }

                Pack.UInt32_To_BE(++ct, buf, 0);
                digest.BlockUpdate(buf, 0, 4);
                digest.DoFinal(buf, 0);

                int xorLen = Math.Min(digestSize, encData.Length - off);
                Xor(encData, buf, off, xorLen);
                off += xorLen;
            }
        }

        private void Xor(byte[] data, byte[] kdfOut, int dOff, int dRemaining)
        {
            for (int i = 0; i != dRemaining; i++)
            {
                data[dOff + i] ^= kdfOut[i];
            }
        }
    }
}
